"""
Модуль для анализа и визуализации данных о кластерах и связях между ними.

Импорт необходимых модулей:
"""
import networkx as nx
import numpy as np
import pandas as pd
from scipy.spatial import distance

def generate_graph(df_objects, df_clusters, base_point = None):
    """
    Генерирует граф сети на основе информации об объектах и кластерах.

    :param df_objects: Сет данных с информацией об объектах.
    :param df_clusters: Сет данных с информацией о кластерах.
    :param base_point: Координаты базовой станции (по умолчанию None).

    :return: Граф сети.
    """
    Graph = nx.Graph()

    Graph = _add_leaders_by_cluster(Graph, df_clusters)

    Graph = _add_base_node_and_edges(Graph, df_objects, base_point, True)

    Graph = _add_objects(Graph, df_objects)

    return Graph

def filter_connections(G, base_connections, leaders_connections):
    """
    Фильтрует граф сети по количеству коннектов между объектами.

    :param G: Граф сети.
    :param base_connections: Количество ближайших базовых коннекторов.
    :param leaders_connections: Количество ближайших лидеров.

    :return: Модифицированный граф сети.
    """
    filtered_G = filter_base_connections(G, "base_connector", base_connections)

    filtered_G = filter_edges_by_group(filtered_G, "leaders", leaders_connections)

    return filtered_G


def short_paths(G, df_objects, target = "Base"):
    """
    Вычисляет кратчайшие пути от каждого узла в графе G до целевого узла "Base".

    Args:
        G (networkx.Graph): Граф, в котором нужно найти кратчайшие пути.
        df_objects (pandas.DataFrame): DataFrame с информацией об объектах.
        target (str, optional): Имя целевого узла. По умолчанию "Base".

    Returns:
        pandas.DataFrame: DataFrame с добавленными столбцами, содержащими кратчайшие пути и длины путей.
    """
    shortest_paths = []  # Список для хранения кратчайших путей
    shortest_path_lengths = []  # Список для хранения длин кратчайших путей
    shortest_path_steps = []  # Список для хранения количества шагов в кратчайших путях
    shortest_path_steps_length = []  # Список для хранения длин каждого шага в кратчайших путях
    id_usage_count = {}  # Словарь для подсчета использования узлов

    for _, row in df_objects.iterrows():
        node = row['object_id']
        shortest_path = nx.shortest_path(G, node, target, weight="weight")
        shortest_path_length = nx.shortest_path_length(G, node, target, weight="weight")
        shortest_path_length = round(shortest_path_length, 2)

        shortest_path_steps.append(len(shortest_path) - 2)

        for id_node in shortest_path[1:-1]:
            id_usage_count[id_node] = id_usage_count.get(id_node, 0) + 1

        shortest_path_steps_length.append([G[shortest_path[j]][shortest_path[j + 1]]['weight'] for j in range(len(shortest_path) - 1)])
        shortest_paths.append(shortest_path)
        shortest_path_lengths.append(shortest_path_length)

    # Добавляем столбцы с результатами в DataFrame
    df_objects['shortest_path'] = shortest_paths
    df_objects['shortest_path_length'] = shortest_path_lengths
    df_objects['shortest_path_steps'] = shortest_path_steps
    df_objects['shortest_path_steps_length'] = shortest_path_steps_length
    df_objects['shortest_path_usage_count'] = df_objects['object_id'].map(id_usage_count).fillna(0).astype(int)

    return df_objects

def filter_base_connections(G, edges_group, connections = 3):
    """
    Фильтрует граф сети по количеству базовых коннекторов.

    :param G: Граф сети.
    :param edges_group: Группа рёбер для фильтрации (например, 'base_connector').
    :param connections: Количество базовых коннекторов для оставления (по умолчанию 3).

    :return: Модифицированный граф сети.
    """
    group_edges = []
    for u, v, g in G.edges(data=True):
        if g.get('group') == edges_group:
            group_edges.append((u, v, g))

    # Сортируем рёбра по весу
    sorted_edges = sorted(group_edges, key=lambda x: x[2]['weight'])
    # Оставляем только connections рёбер
    min_edges = sorted_edges[:connections]

    # Добавляем все остальные рёбра из других групп
    for u, v, g in G.edges(data=True):
        if g.get('group') != edges_group:
            min_edges.append((u, v, g))

    # Создаем граф с отфильтрованными рёбрами
    filtered_G = nx.Graph(min_edges)
    filtered_G.add_nodes_from(G.nodes(data=True))

    return filtered_G

def filter_edges_by_group(G, edges_group, connections = 3):
    """
    Фильтрует рёбра графа по группе.

    :param G: Граф сети.
    :param edges_group: Группа рёбер для фильтрации (например, 'base_connector').
    :param connections: Количество рёбер для оставления (по умолчанию 3).

    :return: Модифицированный граф сети.
    """
    min_edges = []
    for node in G.nodes():
        edges = sorted(G.edges(node, data=True), key=lambda x: x[2]['weight'])
        filtered_edges = [edge for edge in edges if edge[2].get('group') == edges_group]

        min_edges.extend(filtered_edges[:connections])

    # Добавляем все остальные рёбра из других групп
    for u, v, g in G.edges(data=True):
        if g.get('group') != edges_group:
            min_edges.append((u, v, g))

    filtered_G = nx.Graph(min_edges)
    filtered_G.add_nodes_from(G.nodes(data=True))

    return filtered_G


def _add_leaders_by_cluster(G, df_clusters):
    """
    Добавляет лидеров в граф сети на основе информации о кластерах.

    :param G: Граф сети.
    :param df_clusters: Сет данных с информацией о кластерах.

    :return: Модифицированный граф сети.
    """
    for i, cluster in df_clusters.iterrows():
        node_1 = int(cluster['leader_id'])

        G.add_node(node_1, pos=(cluster['leader_x'], cluster['leader_y'], cluster['leader_z']), group="leader")

        leader_coordinates = np.array([cluster['leader_x'], cluster['leader_y'], cluster['leader_z']])

        for j, cluster_2 in df_clusters.iterrows():
            if i != j:
                node_2 = int(cluster_2['leader_id'])

                leader_2_coordinates =  np.array([cluster_2['leader_x'], cluster_2['leader_y'], cluster_2['leader_z']])
                distance_value = distance.euclidean(leader_coordinates, leader_2_coordinates)

                G.add_edge(node_1, node_2, weight=distance_value, group="leaders")

    return G

def _add_base_node_and_edges(G, df_clusters, base_point, filtered = False):
    """
    Добавляет базовую станцию и рёбра к лидерам в граф сети.

    :param G: Граф сети.
    :param df_clusters: Сет данных с информацией о кластерах.
    :param base_point: Координаты базовой станции.
    :param filtered: Флаг фильтрации (по умолчанию False).

    :return: Модифицированный граф сети.
    """
    G.add_node('Base', pos=(base_point), group="base")

    if filtered and 'base_connector' in df_clusters:
        filtered_clusters = df_clusters[df_clusters['base_connector']]
    else:
        filtered_clusters = df_clusters

    for _, cluster in filtered_clusters.iterrows():

        G.add_edge('Base', cluster['leader_id'], weight=cluster['distance_to_base'], group="base_connector")

    return G

def _add_objects(G, df_objects):
    """
    Добавляет объекты в граф сети на основе информации об объектах.

    :param G: Граф сети.
    :param df_objects: Сет данных с информацией об объектах.

    :return: Модифицированный граф сети.
    """
    objects = df_objects.loc[~df_objects['leader']]

    for _, obj in objects.iterrows():
        node = obj['object_id']
        leader_node = obj['leader_id']
        G.add_node(node, pos=(obj['x'], obj['y'], obj['z']), group="object")
        G.add_edge(node, leader_node, weight=obj['distance_to_leader'], group="object")

    return G

def print_stat(G):
    """
    Выводит статистику о графе G.

    Args:
        G (networkx.Graph): Граф.

    Returns:
        None
    """
    edges_list = []
    for edge in G.edges(data=True):
        u, v, data = edge
        weight = data['weight']
        group = data['group']
        coord_u = G.nodes[u]['pos']
        coord_v = G.nodes[v]['pos']
        edges_list.append((u, v, weight, coord_u, coord_v, group))

    df_edges_list = pd.DataFrame(edges_list, columns=['Source', 'Target', 'Weight', 'Source Coords', 'Target Coords', "Group"])

    total_nodes = G.number_of_nodes()
    total_edges = len(G.edges)
    # Получите словарь с количеством связей для каждого узла
    node_degrees = dict(G.degree())
    # Получение общего количества связей
    total_degree_count = sum(node_degrees.values())

    print(f"Общее количество нод: {total_nodes}, рёбер: {total_edges}, связей: {total_degree_count}", end="\n\n")

    print("Статистика дистанций по группам связей графа:", end="\n\n")
    grouped_stats = df_edges_list.groupby('Group')['Weight'].agg(['count', 'min', 'max', 'median', 'mean']).reset_index()

    # Создаём новую строку с общим значением
    total_row = pd.DataFrame({
        'Group': ['Все связи'],
        'count': [df_edges_list['Weight'].count()],
        'min': [df_edges_list['Weight'].min()],
        'max': [df_edges_list['Weight'].max()],
        'median': [df_edges_list['Weight'].median()],
        'mean': [df_edges_list['Weight'].mean()]
    })

    # Добавляем новую строку к grouped_stats
    grouped_stats = pd.concat([grouped_stats, total_row])
    grouped_stats.reset_index(drop=True, inplace=True)
    # Переименование столбцов
    grouped_stats.rename(columns={
        'Group': 'Группа связей',
        'min': 'Минимум',
        'max': 'Максимум',
        'median': 'Медиана',
        'mean': 'Среднее',
        'count': 'Количество ребёр'
    }, inplace=True)

    print(grouped_stats, end="\n\n")
    print("base_connector - cвязь с базовой станцией, object - связь в кластере, leader - связь между кластерами", end="\n\n")
    node_data = []
    for node, attributes in G.nodes(data=True):
        x, y, z = attributes['pos']
        group = attributes['group']
        node_data.append({'node': node, 'x': x, 'y': y, 'z': z, 'group': group})

    df_nodes = pd.DataFrame(node_data)

    group_counts = df_nodes.groupby('group').size()

    df_nodes_groups = pd.DataFrame({'Группа нод': group_counts.index, 'Количество нод': group_counts.values})

    print("Статистика по группам нод графа:", end="\n\n")

    print(df_nodes_groups, end="\n\n")
    print("base - базовая станция, leader - главный в кластере, object - объект", end="\n\n")
    print(f"Общее количество связей: {total_degree_count}", end="\n\n")

    # Создайте Series с количеством узлов для каждого значения degree
    degree_counts = pd.Series(node_degrees).value_counts()
    # Создайте DataFrame из серии
    df_degree_counts = pd.DataFrame({'Количество связей': degree_counts.index, 'Количество нод': degree_counts.values})

    print(df_degree_counts, end="\n\n")

def print_nodes(G):
    """
    Выводит информацию о узлах графа G.

    Args:
        G (networkx.Graph): Граф.

    Returns:
        None
    """
    print( G.nodes(data=True))
    node_data = []
    for node, attributes in G.nodes(data=True):
        x, y, z = attributes['pos']
        group = attributes['group']
        node_data.append({'node': node, 'x': x, 'y': y, 'z': z, 'group': group})

    df_nodes = pd.DataFrame(node_data)

    print(df_nodes)


def routing_stat(G, df_clusters):

    edges_list = []
    for edge in G.edges(data=True):
        u, v, data = edge
        weight = data['weight']
        group = data['group']
        coord_u = G.nodes[u]['pos']
        coord_v = G.nodes[v]['pos']
        edges_list.append((u, v, weight, coord_u, coord_v, group))

    df_edges_list = pd.DataFrame(edges_list, columns=['Source', 'Target', 'Weight', 'Source Coords', 'Target Coords', "Group"])

    leaders = df_edges_list[df_edges_list['Group'] == 'leaders']

    # Создаем словарь для хранения весов лидеров
    data = {}
    for _, row in leaders.iterrows():
        source_id = row['Source']
        target_id = row['Target']
        weight = row['Weight']

        # Если leader_id уже есть в словаре, добавляем вес в список
        if source_id in data:
            data[source_id]['channels_counts'] += 1  # Увеличиваем счетчик весов
            data[source_id]['channels_contacts'].append(target_id)
            data[source_id]['channels_weights'].append(weight)
        else:
            # Иначе создаем новую запись со списком весов
            data[source_id] = {'channels_counts': 1, 
                               'channels_contacts': [target_id],
                               'channels_weights': [weight]}

        # Если leader_id уже есть в словаре, добавляем вес в список
        if target_id in data:
            data[target_id]['channels_counts'] += 1  # Увеличиваем счетчик весов
            data[target_id]['channels_contacts'].append(source_id)
            data[target_id]['channels_weights'].append(weight)
        else:
            # Иначе создаем новую запись со списком весов
            data[target_id] = {'channels_counts': 1, 
                               'channels_contacts': [source_id],
                               'channels_weights': [weight]}

    df_data = pd.DataFrame.from_dict(data, orient='index')
    # Устанавливаем индексную колонку
    df_data.index.name = 'leader_id'

    df_clusters = df_clusters.merge(df_data, left_on='leader_id', right_index=True, how='left')
    
    return df_clusters, df_edges_list
    